
import argparse
import sys
import yaml
from scipy.spatial.transform import Rotation as R
import numpy as np
import json

def validate_args(args):
    if not args.y.endswith('.yaml') and not args.y.endswith('.yml'):
        print("错误: -y 参数必须是一个 YAML 文件")
        sys.exit(1)
    if not args.t.endswith('.json'):
        print("错误: -t 参数必须是一个 JSON 文件")
        sys.exit(1)
    if not args.o.endswith('.yaml') and not args.o.endswith('.yml'):
        print("错误: -o 参数必须是一个 YAML 文件")
        sys.exit(1)

def vec2quat(forward, left):
    """
    forward: list of 3 floats
    left: list of 3 floats
    return: list of 4 floats (x, y, z, w)
    """
    def normalize(v):
        norm = np.linalg.norm(v)
        if norm == 0:
            return v
        return v / norm
    forward = normalize(np.array(forward))
    left = normalize(np.array(left))
    up = normalize(np.cross(forward, left))
    rot_matrix = np.column_stack((left, up, forward))
    rotation = R.from_matrix(rot_matrix)
    quat = rotation.as_quat()  # Returns in (x, y, z, w) format
    return quat.tolist()

def get_pose_from_yaml(yaml_data):
    """
        return:
        {
            "position": [x, y, z],
            "orientation": [x, y, z, w]
        }
    """
    res = dict()
    camera_list = yaml_data['static_camera_list']
    camera = dict()
    for cam in camera_list:
        if cam['name'] == 'midBack_camera':
            camera = cam
            break
    if not camera:
        print("错误: 未找到名为 'midBack_camera' 的相机")
        sys.exit(1)
    res['position'] = camera['position']
    forward = camera['forward']
    left = camera['left']
    res['orientation'] = vec2quat(forward, left)
    return res

def get_transform_matrix(transform_data):
    """
    transform_data: dict loaded from to_adjust.json
    return: 4x4 numpy array
    """
    trans_mat = np.eye(4)
    pos = transform_data['translation']
    pos = [pos['x'], pos['y'], pos['z']]
    rot = transform_data['rotation']
    rot = [rot['roll'], rot['pitch'], rot['yaw']]
    if 'translation' in transform_data:
        trans_mat[:3, 3] = pos
    if 'rotation' in transform_data:
        rot = R.from_euler('xyz', rot)
        trans_mat[:3, :3] = rot.as_matrix()
    return trans_mat

def get_cam(pose, trans_mat):
    """
    pose: dict containing 'position' and 'orientation'
    trans_mat: 4x4 numpy array
    return: dict containing 'position' and 'orientation'
    """
    cam = dict()
    cam['position'] = np.dot(trans_mat, np.array(pose['position'] + [1]))[:3].tolist()
    cam['orientation'] = R.from_matrix(R.from_quat(pose['orientation']).apply(trans_mat[:3, :3])).as_quat().tolist()
    return cam

def get_yaml_data_from_pose(pose, name, sample_yaml_data):
    """
    pose: dict containing 'position' and 'orientation'
    name: str, name of the camera to update
    sample_yaml_data: original yaml data to copy other fields from
    return: new yaml data with updated camera pose
    """
    new_yaml_data = sample_yaml_data.copy()
    camera_list = new_yaml_data['static_camera_list']
    type_name = ''

    for cam in camera_list:
        if cam['name'] == 'midBack_camera':
            type_name = cam['type']
            break
    new_cam = {
        'name': name,
        'type': type_name,
        'position': pose['position'],
        'forward': R.from_quat(pose['orientation']).apply([0, 0, -1]).tolist(),
        'left': R.from_quat(pose['orientation']).apply([-1, 0, 0]).tolist()
    }
    new_yaml_data['static_camera_list'].append(new_cam)
    return new_yaml_data

def main():
    usage="""
用法: python cam_yaml.py -y <yaml_file> -t <transform_file> -o <output_file>
说明: 只允许 -y、-t 和 -o 三个参数
    如果有 -o 则输出到指定文件，否则覆盖 -y 文件
"""
    parser = argparse.ArgumentParser(description=usage)
    parser.add_argument('-y', required=True, help='yaml 文件路径')
    parser.add_argument('-t', required=True, help='to_adjust.json 文件路径')
    parser.add_argument('-o', required=False, help='输出路径')
    args, unknown = parser.parse_known_args()
    if unknown:
        print(f"不支持的参数: {' '.join(unknown)}")
        sys.exit(1)
    validate_args(args)

    yaml_path = args.y
    with open(yaml_path, 'r', encoding='utf-8') as f:
        yaml_data = yaml.safe_load(f)
    midBack_pose = get_pose_from_yaml(yaml_data)
    transform_path = args.t
    with open(transform_path, 'r', encoding='utf-8') as f:
        transform_data = json.load(f)
    trans_mat = get_transform_matrix(transform_data)
    midBack_real_camera_pose = get_cam(midBack_pose, trans_mat)
    midBack_real_inv_camera_pose = get_cam(midBack_pose, np.linalg.inv(trans_mat))
    new_yaml_data = get_yaml_data_from_pose(midBack_real_inv_camera_pose, 'midBack_real_inv_camera', yaml_data)
    new_yaml_data = get_yaml_data_from_pose(midBack_real_camera_pose, 'midBack_real_camera', new_yaml_data)

    output_path = args.o if args.o else yaml_path
    with open(output_path, 'w', encoding='utf-8') as f:
        yaml.safe_dump(new_yaml_data, f, allow_unicode=True, sort_keys=False)

if __name__ == "__main__":
    main()
